#include <iostream>
#include <string>
#include <ctime>
#include <sstream>
#include <iomanip>
#include <winsock2.h>
#include <ws2tcpip.h>
using namespace std;

#pragma comment(lib, "ws2_32.lib")

//NTP协议相关定义
#define NTP_PORT 123
#define NTP_PACKET_SIZE 48
#define NTP_TIMESTAMP_DELTA 2208988800ull // 1900-1970 秒数差

//将时间字符串转换为NTP时间戳
uint64_t time_string_to_ntp(const string& time_str)
{
    struct tm tm = {0};
    istringstream ss(time_str);
    
    //解析时间字符串，格式为"YYYY-MM-DD HH:MM:SS"
    ss >> get_time(&tm, "%Y-%m-%d %H:%M:%S");
    if (ss.fail()) throw runtime_error("Failed to parse time string");
    
    //转换为 time_t (UTC时间)
    time_t t = _mkgmtime(&tm);
    if (t == -1) throw runtime_error("Invalid time");
    
    //转换为NTP时间戳(从1900年开始的秒数)
    return static_cast<uint64_t>(t) + NTP_TIMESTAMP_DELTA;
}

//将NTP时间戳写入缓冲区
void write_ntp_timestamp(uint8_t* buffer, uint64_t timestamp)
{
    //NTP时间戳是64位定点数，前32位是秒，后32位是小数秒
    uint32_t seconds = static_cast<uint32_t>(timestamp);
    uint32_t fraction = static_cast<uint32_t>((timestamp - seconds) * 0x100000000);
    
    //网络字节序 (大端)
    buffer[0] = (seconds >> 24) & 0xFF;
    buffer[1] = (seconds >> 16) & 0xFF;
    buffer[2] = (seconds >> 8) & 0xFF;
    buffer[3] = seconds & 0xFF;
    
    buffer[4] = (fraction >> 24) & 0xFF;
    buffer[5] = (fraction >> 16) & 0xFF;
    buffer[6] = (fraction >> 8) & 0xFF;
    buffer[7] = fraction & 0xFF;
}

int main(int argc, char* argv[])
{
    if (argc != 2)
	{
        cerr << "Usage: " << argv[0] << " \"YYYY-MM-DD HH:MM:SS\"" << endl;
        return 1;
    }

    string time_str = argv[1];
    uint64_t ntp_timestamp;
    
    try
	{
        ntp_timestamp = time_string_to_ntp(time_str);
    } 
	catch (const exception& e)
	{
        cerr << "Error: " << e.what() << endl;
        return 1;
    }

    //初始化Winsock
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
	{
        cerr << "WSAStartup failed" << endl;
        return 1;
    }

    //创建UDP套接字
    SOCKET sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET)
	{
        cerr << "Socket creation failed: " << WSAGetLastError() << endl;
        WSACleanup();
        return 1;
    }

    //绑定到NTP端口
    sockaddr_in serverAddr;
    serverAddr.sin_family = AF_INET;
    serverAddr.sin_port = htons(NTP_PORT);
    serverAddr.sin_addr.s_addr = htonl(INADDR_ANY);

    if (bind(sock, (sockaddr*)&serverAddr, sizeof(serverAddr)) == SOCKET_ERROR)
	{
        cerr << "Bind failed: " << WSAGetLastError() << endl;
        closesocket(sock);
        WSACleanup();
        return 1;
    }

    cout << "NTP server running. Using time: " << time_str << endl;

    //主循环
    while (true)
	{
        sockaddr_in clientAddr;
        int clientAddrLen = sizeof(clientAddr);
        uint8_t buffer[NTP_PACKET_SIZE] = {0};

        //接收NTP请求
        int bytesReceived = recvfrom(sock, (char*)buffer, NTP_PACKET_SIZE, 0, (sockaddr*)&clientAddr, &clientAddrLen);
        if (bytesReceived == SOCKET_ERROR)
		{
            cerr << "recvfrom failed: " << WSAGetLastError() << endl;
            continue;
        }

        //确保这是一个NTP请求(至少48字节)
        if (bytesReceived < NTP_PACKET_SIZE) continue;

        //设置NTP响应
        buffer[0] = 0x1C;// LI=0, VN=3, Mode=4 (服务器)

        //写入传输时间戳(当前时间)
        write_ntp_timestamp(&buffer[32], ntp_timestamp);
        //写入参考时间戳(同样使用当前时间)
        write_ntp_timestamp(&buffer[16], ntp_timestamp);
        //写入原始时间戳(来自请求)
        memcpy(&buffer[24], &buffer[40], 8);

        //发送响应
        if (sendto(sock, (char*)buffer, NTP_PACKET_SIZE, 0, (sockaddr*)&clientAddr, clientAddrLen) == SOCKET_ERROR) std::cerr << "sendto failed: " << WSAGetLastError() << std::endl;
    }

    //清理
    closesocket(sock);
    WSACleanup();
    return 0;
}